import asyncio
import json
from contextlib import AsyncExitStack
from typing import Any

from mcp import ClientSession, StdioServerParameters, stdio_client, ListToolsResult, ListResourcesResult
from mcp.client.sse import sse_client
from mcp.types import CallToolResult, ReadResourceResult, GetPromptResult

from pre_lesson.chapter8_cclient.mcp_model import MCPFunction


#连接服务端参数信息及状态信息
class ServerInfo:
    def __init__(self,
                 name: str,
                 cmd: str | None = None,
                 args: list[str] | None = None,
                 env: dict[str, Any] | None = None,
                 url: str | None = None,
                 ):
        self.name = name
        if cmd in ['npx', 'uvx', 'node']:
            self.transport_type = 'stdio'
        if url is not None:
            self.transport_type = 'sse'
        self.cmd = cmd
        self.args = args
        self.env = env
        self.url = url

        # 与服务器连接的session对象
        self.session: ClientSession | None = None


class ServerManager:

    def __init__(self):
        '''初始化信息'''
        self.servers: dict[str, ServerInfo] = {}
        self.async_exit_stack = AsyncExitStack()
        self.mcp_functions: dict[str, MCPFunction] = {}

    async def initialize(self):
        '''初始化'''
        await self.load_mcp_json()
        await self.create_session()
        await self.fetch_all_function_calling()
        print("mcp_functions:",self.mcp_functions)

    async def load_mcp_json(self):
        '''加载mcp.json文件'''
        with open('mcp.json', 'r', encoding='utf-8') as file:
          data = json.load(file)
          print("mcp.json 加载成功")
          mcp_servers = data.get("mcpServers")
          #解析mcp.json文件，并初始成成ServerInfo对象
          for key, item in mcp_servers.items():
            server_info = ServerInfo(
              name=key,
              cmd=item.get("command"),
              args=item.get("args"),
              env=item.get("env"),
              url=item.get("url")
            )
            #key = server_name
            self.servers[key] = server_info
    #创建连接会话信息
    async def create_session(self):
        for key, server_info in self.servers.items():
            print("server_name ==> ",key)
            print("server_info ==> ",server_info)
            client = None
            #transport类型
            transport_type = server_info.transport_type
            if transport_type == 'stdio':
                #创建StdioServerParameter参数信息
                server_parameters = StdioServerParameters(
                    command=server_info.cmd,
                    args=server_info.args,
                    env=server_info.env
                )
                #创建stdioClient
                client = stdio_client(server=server_parameters)

            elif transport_type =='sse':
                # 创建sseClient
                client = sse_client(url=server_info.url)

            else:
                print("不支持的transport类型")

            stdio_transport = await self.async_exit_stack.enter_async_context(client)
            read_stream, write_stream = stdio_transport
            # 创建ClientSession
            client_session = ClientSession(read_stream, write_stream)
            # 保存session
            server_info.session = await self.async_exit_stack.enter_async_context(client_session)
            # 初始化session
            await server_info.session.initialize()
            print("session:",server_info.session,"初始化完成")

    async def fetch_all_function_calling(self):
        for key, server_info in self.servers.items():
            session = server_info.session
            print("=== server_name:",key)
            # 列出所有可用的工具
            try:
                list_tool: ListToolsResult = await session.list_tools()
                tools = list_tool.tools
            except Exception as e:
                print("list_tools error:",e)
                tools = []
            print("================== wrapper tool ======================")
            for tool in tools:
                print("tool ==> ", tool.name, ",", tool.description, ",", tool.inputSchema)
                #初始化MCPFunction
                mcp_function = MCPFunction(
                    name=tool.name,
                    function_name=tool.name,
                    function_type='tool',
                    description=tool.description,
                    input_schema=tool.inputSchema,
                    server_name=key
                )
                #保存到mcp_functions
                self.mcp_functions[tool.name] = mcp_function
            print("================== wrapper resource ======================")
            # 列出所有的资源列表
            try:
                list_resources_result: ListResourcesResult = await session.list_resources()
                resources = list_resources_result.resources
            except Exception as e:
                print("list_resources error:",e)
                resources = []
            for resource in resources:
                #特殊处理name名称有空格的情况
                resource.name = resource.name.replace(" ", "")
                print("resource ==> ", resource.name, ",", resource.uri, ",", resource.description, ",",
                      resource.mimeType)

                #初始化MCPFunction
                mcp_function = MCPFunction(
                    name=resource.name,
                    server_name=key,
                    function_name=resource.name,
                    function_type='resource',
                    description=resource.description,
                    uri=resource.uri,
                    arguments=None
                )
                #保存到mcp_functions
                self.mcp_functions[resource.name] = mcp_function

            # 列表所有提示类
            try:
                list_prompts_result = await session.list_prompts()
                prompts = list_prompts_result.prompts
            except Exception as e:
                print("list_prompts error:",e)
                prompts = []
            print("================== wrapper prompt ======================")
            for prompt in prompts:
                print("prompt ==> ", prompt.name, ",", prompt.description, ",", prompt.arguments)
                #初始化MCPFunction
                mcp_function = MCPFunction(
                    function_name=prompt.name,
                    name=prompt.name,
                    function_type='prompt',
                    description=prompt.description,
                    arguments=prompt.arguments,
                    server_name=key
                )
                #保存到mcp_functions
                self.mcp_functions[prompt.name] = mcp_function
    # function calling 方法执行
    async def execute_calling(self,function_name:str,arguments:dict[str,Any]):
        mcp_function = self.mcp_functions[function_name]
        function_type = mcp_function.function_type
        session = self.servers[mcp_function.server_name].session
        if function_type == 'tool':
            tool_result:CallToolResult = await session.call_tool(name=function_name, arguments=arguments)
            print("====  调用工具链结果：",tool_result)
            return tool_result.content[0].text
        elif function_type == 'resource':
            uri = mcp_function.uri
            resource_result:ReadResourceResult = await session.read_resource(uri)
            print("====  调用资源链结果：",resource_result)
            return resource_result.contents[0].text
        elif function_type =='prompt':
            prompt_result:GetPromptResult = await session.get_prompt(name=function_name, arguments=arguments)
            print("====  调用prompt链结果：",prompt_result)
            return prompt_result.messages[0].content.text
        else:
            print("========不支持的function_type============")


    async def aclose(self):
        await self.async_exit_stack.aclose()

async def test():
    server_manager = ServerManager()
    try:

        await server_manager.initialize()
        #await server_manager.execute_calling("read_file",{"path":"D:\dir_test\\test3.txt"})
        #await server_manager.execute_calling("list_tables",{})
        result =await server_manager.execute_calling("mcp-demo",{"topic": "mcp-demo-topic"})
        print("result:",result)
    except Exception as e:
        print("Error:",e)

    finally:
        await server_manager.async_exit_stack.aclose()

if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(test())

